/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/**
 * @file channel_mixing_310b.h
 */
#ifndef CHANNEL_MIXING_310B_H
#define CHANNEL_MIXING_310B_H
#include "kernel_operator.h"
#include "lib/matmul_intf.h"
using namespace AscendC;
using namespace matmul;
//ub:253952

class ChannelMixingKernel310b {
public:
    __aicore__ inline ChannelMixingKernel310b() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR h0, GM_ADDR xk, GM_ADDR kw, GM_ADDR vw, GM_ADDR out, GM_ADDR ht, 
                                GM_ADDR hs, const TCubeTiling &tiling1, const TCubeTiling &tiling2, 
                                int32_t B, int32_t T,
                                TPipe *pipe) 
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->tiling1 = tiling1;
        this->tiling2 = tiling2;
        this->T = T;

        xGm.SetGlobalBuffer((__gm__ half*)x + B*T*C);// - this->T*C);
        h0Gm.SetGlobalBuffer((__gm__ half*)h0 + B*C);// - C);
        xkGm.SetGlobalBuffer((__gm__ half*)xk);
        kwGm.SetGlobalBuffer((__gm__ half*)kw);
        vwGm.SetGlobalBuffer((__gm__ half*)vw);

        outGm.SetGlobalBuffer((__gm__ half*)out + B*T*C);// - this->T*C);
        htGm.SetGlobalBuffer((__gm__ half*)ht + B*C);// - C);

        hsGm.SetGlobalBuffer((__gm__ half*)hs);


        pipe->InitBuffer(inQueueX, BUFFER_NUM, C * 2 * sizeof(half));
        pipe->InitBuffer(inQueueXk, 1, C * sizeof(half));

        pipe->InitBuffer(outQueueOut, BUFFER_NUM, C * sizeof(half));

        if (GetSysWorkSpacePtr() == nullptr) {
            return;
        }
    }
    
    __aicore__ inline void Process(TPipe *pipe)
    {
      //token_shift
      // decode T=1
      CopyIn1Rwk();
      CopyIn1H0();
      Compute1(0);
      CopyOut1(0);

      // prefill T>1
      for (int32_t i = 0; i < this->T-1; i++) 
      {
          CopyIn1(i);
          Compute1(i);
          CopyOut1(i+1);
      }
      LocalTensor<half> xkLocal = inQueueXk.DeQue<half>();
      inQueueXk.FreeTensor(xkLocal);

      {
        LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
        DataCopy(xLocal, xGm[(this->T-1) * C], C);
        inQueueX.EnQue(xLocal);
        xLocal = inQueueX.DeQue<half>();
        LocalTensor<half> outLocal = outQueueOut.AllocTensor<half>();
        Adds(outLocal, xLocal, (half)0, C);
        outQueueOut.EnQue<half>(outLocal);
        inQueueX.FreeTensor(xLocal);
        outLocal = outQueueOut.DeQue<half>();  
        DataCopy(htGm, outLocal, C);
        outQueueOut.FreeTensor(outLocal);
      }

      ///////////////////////

      matmulObj.Init(&this->tiling1);
      matmulObj.SetTensorA(outGm);
      matmulObj.SetTensorB(kwGm, true);
      matmulObj.IterateAll(hsGm);
      matmulObj.End();

      for (int32_t i = 0; i < this->T*4; i++) 
      {
          CopyIn2(i);
          Compute2(i);
          CopyOut2(i);
      }

      matmulObj.Init(&this->tiling2);
      matmulObj.SetTensorA(hsGm);
      matmulObj.SetTensorB(vwGm, true);
      matmulObj.IterateAll(outGm);
      matmulObj.End();

      PipeBarrier<PIPE_ALL>();
    }

private:
    __aicore__ inline void CopyIn1Rwk()
    {
      LocalTensor<half> xkLocal = inQueueXk.AllocTensor<half>();
      DataCopy(xkLocal, xkGm, C);
      inQueueXk.EnQue(xkLocal);
    }
    __aicore__ inline void CopyIn1H0()
    {
      LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
      DataCopy(xLocal, h0Gm, C);
      DataCopy(xLocal[C], xGm, C);
      inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void CopyIn1(int32_t progress)
    {
      LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
      DataCopy(xLocal, xGm[progress * C], C * 2);
      inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute1(int32_t progress)
    {
      LocalTensor<half> xLocal = inQueueX.DeQue<half>();
      LocalTensor<half> xkLocal = inQueueXk.DeQue<half>();
      LocalTensor<half> outLocal = outQueueOut.AllocTensor<half>();

      Sub(outLocal, xLocal, xLocal[C], C);
      Mul(outLocal, outLocal, xkLocal, C);
      Add(outLocal, outLocal, xLocal[C], C);

      outQueueOut.EnQue<half>(outLocal);
      inQueueX.FreeTensor(xLocal);
      inQueueXk.EnQue(xkLocal);
    }
    
    __aicore__ inline void CopyOut1(int32_t progress)
    {
      LocalTensor<half> outLocal = outQueueOut.DeQue<half>();  
      DataCopy(outGm[progress * C], outLocal, C);
      outQueueOut.FreeTensor(outLocal);
    }
    ////////////
    __aicore__ inline void CopyIn2(int32_t progress)
    {
      LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
      DataCopy(xLocal, hsGm[progress * C], C);
      inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute2(int32_t progress)
    {
      LocalTensor<half> xLocal = inQueueX.DeQue<half>();
      LocalTensor<half> outLocal = outQueueOut.AllocTensor<half>();


      Relu(xLocal, xLocal, C);
      Mul(outLocal, xLocal, xLocal, C);
      
      outQueueOut.EnQue<half>(outLocal);
      inQueueX.FreeTensor(xLocal);
    }
    
    __aicore__ inline void CopyOut2(int32_t progress)
    {
      LocalTensor<half> outLocal = outQueueOut.DeQue<half>();  
      DataCopy(hsGm[progress * C], outLocal, C);
      outQueueOut.FreeTensor(outLocal);
    }
    __aicore__ inline uint32_t Ceiling(uint32_t a, uint32_t b)
    {
        return (a + b - 1) / b;
    }
public:
    static constexpr int32_t BUFFER_NUM = 2;
    static constexpr int32_t C = 2560;
    static constexpr MatmulConfig MM_CFG = GetNormalConfig();//(false, false, true);
    Matmul<MatmulType<TPosition::GM, CubeFormat::ND, half>, 
           MatmulType<TPosition::GM, CubeFormat::ND, half, true>,
           MatmulType<TPosition::GM, CubeFormat::ND, half>, 
           MatmulType<TPosition::GM, CubeFormat::ND, half>,
           MM_CFG>
        matmulObj;

    GlobalTensor<half> xGm;
    GlobalTensor<half> h0Gm;
    GlobalTensor<half> xkGm;
    GlobalTensor<half> kwGm;
    GlobalTensor<half> vwGm;
    GlobalTensor<half> outGm;
    GlobalTensor<half> htGm;

    GlobalTensor<half> hsGm;

    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECIN, 1> inQueueXk;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOut;


    int32_t T;

    TCubeTiling tiling1;
    TCubeTiling tiling2;
};
#endif // CHANNEL_MIXING_310B_H
